package com.mgq.algorithm.unionfind;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集实现
 *
 */
public class UnionFind<V> {
    //elementMap 保存的是V和Element的对应关系
    private HashMap<V, Element<V>> elementMap;
    //fatherMap保存的是Element的father(head)元素是谁 (相当于一个向上的指针)(链表也可以,此处使用map来表示)
    private HashMap<Element<V>, Element<V>> fatherMap;
    //sizeMap中保存的是代表元素的大小(是整个集合head包含的节点个数)
    private HashMap<Element<V>, Integer> sizeMap = new HashMap<>();

    /**
     * 并查集在初始化时,要求用户必须将样本全部给我们
     * @param list
     */
    public UnionFind(List<V> list) {
        elementMap = new HashMap<>();
        fatherMap = new HashMap<>();
        for (V v : list) {
            Element<V> element = new Element<>(v);
            elementMap.put(v, element);
            fatherMap.put(element, element);
            //一开始只包含自己
            sizeMap.put(element, 1);
        }
    }

    /**
     * 查找元素的head
     *
     * @param element
     * @return
     */
    public Element<V> findHead(Element<V> element) {
        Stack<Element<V>> stack = new Stack<>();
        while (fatherMap.get(element) != element) {
            stack.push(element);
            element = fatherMap.get(element);
        }
        //扁平化结构, 就是把检查过的元素,直接指向head,这样下次查询的时候时间复杂度直接变为O(1)
        //相当于是优化了结构
        while (!stack.isEmpty()) {
            fatherMap.put(stack.pop(), element);
        }
        return element;
    }

    /**
     * a,b是否在同一个集合
     *
     * @param a
     * @param b
     * @return
     */
    public boolean isSameSet(V a, V b) {
        if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
            return false;
        }
        Element<V> aH = findHead(elementMap.get(a));
        Element<V> bH = findHead(elementMap.get(b));
        return aH == bH;
    }

    /**
     * 合并2个集合
     *
     * @param a
     * @param b
     */
    public void union(V a, V b) {
        if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
            Element<V> aF = findHead(elementMap.get(a));
            Element<V> bF = findHead(elementMap.get(b));
            if (aF != bF) {
                Element<V> big = sizeMap.get(elementMap.get(a)) > sizeMap.get(elementMap.get(b)) ? aF : bF;
                Element<V> small = big == aF ? bF : aF;
                //把小的集合合并到大的集合中
                fatherMap.put(small,big);
                sizeMap.put(big, sizeMap.get(big) + sizeMap.get(small));
                sizeMap.remove(small);
            }
        }
    }


    /**
     * Element相当于把V包装了一层
     *
     * @param <V>
     */
    static class Element<V> {
        private V value;

        public Element(V value) {
            this.value = value;
        }
    }
}
